import cv2
from mmocr.utils.ocr import MMOCR

class PT_Rec:
    def __init__(self, rec_config, rec_checkpoint):
        self.recog_model = MMOCR(det=None,
                        recog='SATRN',
                        recog_config=rec_config,
                        recog_ckpt=rec_checkpoint)
        

    def __call__(self, im):
        '''

        :param im: numpy
        :return: string
        '''
        recog_text = self.recog_model.readtext(im)[0]['text']
        return recog_text


if __name__=='__main__':
    rec_config = 'configs/textrecog/satrn/satrn_pt.py'
    rec_checkpoint = 'epoch_6.pth'
    model=PT_Rec(rec_config, rec_checkpoint)
    im=cv2.imread('1.jpg')
    res=model(im)